import json
import time
import traceback

from flask import Flask, request, jsonify
import cv2
import numpy as np
import base64

from test_seg_kakou import get_algo_fubahanxi, get_algo_hanqianxianshu
from global_vars import logger, profiler, isDebug

#
engine_dict = {
    'hanhou': get_algo_fubahanxi(),
    'hanqianxianshu': get_algo_hanqianxianshu()
}


app = Flask(__name__)

@app.route('/')
def index():
    return 'hello world'

@app.route('/detect', methods=['POST'])
@profiler
def detect():
    # return jsonify({'res': 0})  # to respose
    # return json.dumps({'res': 0})
    try:
        # 1 获取 base64 编码的图片数据
        logger.info(f"1 获取请求数据")
        data = request.get_json()
        if 'image_base64' not in data:
            return jsonify({'error': 'No image data provided'}), 400
        image_base64 = data['image_base64']
        funct_str = data['function'] if data.get('function') else 'hanhou'
        product = data['product'] if data.get('product') else '519'
        logger.info(f"已获取请求数据 base64图片: {len(image_base64)}; 算法: {funct_str}; 产品型号: {product}")

        image_array = np.frombuffer(base64.b64decode(image_base64), dtype=np.uint8)
        # 解码 base64 数据
        # 将图片数组解码为 OpenCV 图像
        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
        # cv2.imwrite(r'D:\code\git\zxc\project_fubahanxi\Pictures\1.jpg', image)

        # 2 图像处理逻辑
        logger.info("2 图像处理逻辑")
        result = process_image(image, funct_str=funct_str, product=product)
        # 3 返回结果
        logger.info(f"3 返回结果 {result['result']}")
        if isDebug:
            logger.info(profiler.print_stats()) # 单位1/10000ms 累加
            # profiler.clear_stats()
        return jsonify(result) # to respose
    except:
        error_str = traceback.format_exc()
        logger.error(error_str)
        return jsonify({'error': error_str}), 400



@profiler
def process_image(image, funct_str='hanhou', product='519'):
    # 在这里添加你的图像处理逻辑
    # 例如，我们可以简单地返回图像的形状作为检测结果
    show_img, out, end, ends, save_img = engine_dict[funct_str](image, product)
    if isDebug:
        cv2.imwrite(rf'Pictures\server\{time.time()}.jpg', show_img)
    show_img_base64 = base64.b64encode(cv2.imencode('.jpg', show_img)[1].tobytes()).decode()  # numpy to base64_st

    res_dict = {}
    if funct_str == 'hanhou':

        res_dict = {
            'function': funct_str,
            'product': product,
            'show_img_base64': show_img_base64,
            'result': end,
            'end': end,
            'massage': '消息',
        }
    elif funct_str == 'hanqianxianshu':
        template_names, pre_clsnames, msg = ends
        res_dict = {
            'function': funct_str,
            'product': product,
            'show_img_base64': show_img_base64,
            'result': end,
            'massage': msg,
            'other': {
                'template_names': template_names,
                'predict_names': pre_clsnames,
                'isoverlap': False,
            }

        }

    return res_dict


if __name__ == '__main__':
    # app.run(debug=True)
    app.run(host='0.0.0.0', port=5000, debug=False)
